package cn.bugstack.dev.tech.test;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;
import com.alibaba.fastjson.JSON;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @Description: rag测试类
 * @ClassName: RagTest
 * @Author: zhaiyongxin
 * @Date: 2025/7/8 10:37
 * @Version: 1.0
 */
@SpringBootTest
@Slf4j
@RunWith(SpringRunner.class)
public class RagTest {

    @Resource
    private TokenTextSplitter tokenTextSplitter;
    @Resource
    private PgVectorStore pgVectorStore;
    @Resource
    private OllamaChatClient ollamaChatClient;

    @Test
    public void upload() {
        TikaDocumentReader reader = new TikaDocumentReader("./data/file.text");
        List<Document> documents = reader.get();
        List<Document> documentSplitterList  = tokenTextSplitter.apply(documents);

        documents.forEach(doc -> {
            doc.getMetadata().put("knowledge", "知识库");
        });

        documentSplitterList.forEach(doc -> {
            doc.getMetadata().put("knowledge", "知识库名称1");
        });
        pgVectorStore.accept(documentSplitterList);

    }

    @Test
    public void chat() {
        String message = "翟永鑫今年多大";
        String SYSTEM_PROMPT = """
            Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
            If unsure, simply state that you don't know.
            Another thing you need to note is that your reply must be in Chinese!
            DOCUMENTS:
                {documents}
            """;

        //从向量库获取内容
        SearchRequest searchRequest = SearchRequest.query(message)
                .withTopK(5)
                .withFilterExpression("knowledge == '知识库名称1'");
        List<Document> documents = pgVectorStore.similaritySearch(searchRequest);

        List<Message> messages = new ArrayList<>();
        messages.add(new UserMessage(message));
        Message sysMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documents));
        messages.add(sysMessage);
        ChatResponse chatResponse = ollamaChatClient.call(
                new Prompt(
                        messages,
                        OllamaOptions.create()
                                .withModel("deepseek-r1:1.5b")
                ));
        log.info("测试结果：{}", JSON.toJSONString(chatResponse));
    }

}
